__doc__ =   """ Classes to determine simple atom depth.
            """

from itertools import izip
from numpy import sqrt, mean, squeeze, array 

from AbstractPropertyDict import AbstractAtomPropertyDict
from AbstractPropertyDict import AbstractResiduePropertyDict
from AtomProperties import get_united_radius, get_vdw_radius
from Einput import einput
from Externals import call_msms
from Crunchers import kdtree


class ReducedSurface(object):
    """ Class to hold the reduced surface of an atom-based entity.
    """
    def __init__(self, entities, **kwargs):
        atoms = einput(entities, 'A', 'atoms')      # holder of atoms
        self.surface = call_msms(atoms, **kwargs)   # surface array from msms
        self.kdt = kdtree(self.surface)         # kdtree from array / ANN

    def depth(self, entity):
        (indice, depth) = self.kdt.knn(entity.get_coords(), 1)  # find one
        return sqrt(depth)  # return squared distance

    def depths(self, entities):
        depths = []
        coords = (entity.get_coords() for entity in entities)   # find multiple
        for coord in coords:
            (ind, depth) = self.kdt.knn(coord, 1)            
            depths.extend(depth)
        return sqrt(array(depths)) # numpy square on whole array


class AbstractAtomDepth(AbstractAtomPropertyDict):
    """ Returns depth of atoms for given surface.
    """
    def __init__(self, atoms, xtra_key, surface):
        depths = list(squeeze(surface.depths(atoms)))
        result = dict(zip(atoms.keys(), depths))
        for (atom, depth) in izip(atoms.itervalues(), depths):
            atom.xtra[xtra_key] = depth
        AbstractAtomPropertyDict.__init__(self, result)


class AtomDepth(AbstractAtomDepth):
    def __init__(self, entities, surface =None, xtra_key =None, **kwargs):
        xtra_key = (xtra_key or 'DPT')
        atoms = einput(entities, 'A', 'atoms')
        if not surface:
            surface = ReducedSurface(atoms)
        AbstractAtomDepth.__init__(self, atoms, xtra_key, surface)


class BoolAtomDepth(AbstractAtomPropertyDict):
    """ Applies an atom depth threshold to filter atoms close to the surface.
        The default filters is the united atom radius each atom type.
    """
    def __init__(self, entities, srf_key =None, dpt_key =None, max_err =None, r_type =None):
        dpt_key = (dpt_key or 'bDPT')         # default depth
        xtra_key = 'BOOL_' + dpt_key                #
        r_type = (r_type or 'get_united_radius')    # atomic radius
        max_err = (max_err or 0.0)                  # by default msms reduced
        atoms = einput(entities, 'A', 'atoms')
        result = {}
        for (id_, atom) in atoms.iteritems():
            xtra_key_value = False
            dpt_max = eval(r_type)(atom.get_dict()) + max_err   # get the atomic radius
            if atom.xtra[dpt_key] <= dpt_max:
                xtra_key_value = True
            atom.xtra[xtra_key] = xtra_key_value
            result[id_] = xtra_key_value
        AbstractAtomPropertyDict.__init__(self, result)

#EOF